/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.connector.jdbc.table;

import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.connector.jdbc.converter.JdbcRowConverter;
import org.apache.flink.connector.jdbc.dialect.JdbcDialect;
import org.apache.flink.connector.jdbc.internal.connection.JdbcConnectionProvider;
import org.apache.flink.connector.jdbc.internal.connection.SimpleJdbcConnectionProvider;
import org.apache.flink.connector.jdbc.internal.options.InternalJdbcConnectionOptions;
import org.apache.flink.connector.jdbc.statement.FieldNamedPreparedStatement;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.functions.FunctionContext;
import org.apache.flink.table.functions.LookupFunction;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;

/** A lookup function for {@link JdbcDynamicTableSource}. */
@Internal
public class JdbcRowDataLookupFunction extends LookupFunction {

    private static final Logger LOG = LoggerFactory.getLogger(JdbcRowDataLookupFunction.class);
    private static final long serialVersionUID = 2L;

    private final String query;
    
    private final JdbcConnectionProvider connectionProvider;
    
    private final String[] keyNames;
    
    private final int maxRetryTimes;
    
    private final JdbcRowConverter jdbcRowConverter;
    
    private final JdbcRowConverter lookupKeyRowConverter;

    private transient FieldNamedPreparedStatement statement;

    public JdbcRowDataLookupFunction(
            InternalJdbcConnectionOptions options,
            int maxRetryTimes,
            String[] fieldNames,
            DataType[] fieldTypes,
            String[] keyNames,
            RowType rowType) {
        
        checkNotNull(options, "No JdbcOptions supplied.");
        checkNotNull(fieldNames, "No fieldNames supplied.");
        checkNotNull(fieldTypes, "No fieldTypes supplied.");
        checkNotNull(keyNames, "No keyNames supplied.");
        
        this.connectionProvider = new SimpleJdbcConnectionProvider(options);
        
        this.keyNames = keyNames;
        
        List<String> nameList = Arrays.asList(fieldNames);
        
        DataType[] keyTypes = Arrays.stream(keyNames)
                        .map(
                                s -> {
                                    checkArgument(
                                            nameList.contains(s),
                                            "keyName %s can't find in fieldNames %s.",
                                            s,
                                            nameList);
                                    return fieldTypes[nameList.indexOf(s)];
                                })
                        .toArray(DataType[]::new);
        
        this.maxRetryTimes = maxRetryTimes;
        
        this.query = options.getDialect()
                        .getSelectFromStatement(options.getTableName(), fieldNames, keyNames);
        
        JdbcDialect jdbcDialect = options.getDialect();

        this.jdbcRowConverter = jdbcDialect.getRowConverter(rowType);

        this.lookupKeyRowConverter = jdbcDialect.getRowConverter(
                        RowType.of(
                                Arrays.stream(keyTypes)
                                        .map(DataType::getLogicalType)
                                        .toArray(LogicalType[]::new)
                        )
        );
    }

    @Override
    public void open(FunctionContext context) throws Exception {
        try {
            establishConnectionAndStatement();
        } catch (SQLException sqe) {
            throw new IllegalArgumentException("open() failed.", sqe);
        } catch (ClassNotFoundException cnfe) {
            throw new IllegalArgumentException("JDBC driver class not found.", cnfe);
        }
    }

    /**
     * This is a lookup method which is called by Flink framework in runtime.
     * 这是 Flink 框架在运行时调用的查找方法。
     * 
     * 这段代码实现了一个 JDBC数据库查询操作，根据给定的 keyRow查找匹配的行数据，
     * 并将结果作为 RowData集合返回。
     * 
     * @param keyRow lookup keys（查找键）
     */
    @Override
    public Collection<RowData> lookup(RowData keyRow) {
        
        //使用for循环，最多重试 maxRetryTimes次查询操作。
        for (int retry = 0; retry <= maxRetryTimes; retry++) {
            try {
                statement.clearParameters();
                // 在循环中，首先通过 lookupKeyRowConverter.toExternal(keyRow, statement)方法
                // 将 keyRow转换为 JDBC语句中的参数，并将其设置到 statement 对象中。
                statement = lookupKeyRowConverter.toExternal(keyRow, statement);
                
                // 然后执行查询操作，
                try (ResultSet resultSet = statement.executeQuery()) {
                    
                    ArrayList<RowData> rows = new ArrayList<>();
                    
                    while (resultSet.next()) {
                        //并将查询结果转换为 RowData对象，
                        RowData row = jdbcRowConverter.toInternal(resultSet);
                        
                        //并存储在一个 ArrayList中。
                        rows.add(row);
                    }
                    // trimToSize()方法用于将 ArrayList对象的容量调整为其当前大小，
                    // 以便释放未使用的内存。
                    rows.trimToSize();
                    //最后返回 ArrayList中的所有行数据。
                    return rows;
                }
            } catch (SQLException e) {
                LOG.error(String.format("JDBC executeBatch error, retry times = %d", retry), e);
                if (retry >= maxRetryTimes) {
                    throw new RuntimeException("Execution of JDBC statement failed.", e);
                }

                try {
                    // 在重试过程中，如果连接无效，
                    if (!connectionProvider.isConnectionValid()) {
                        // 则关闭连接
                        statement.close();
                        connectionProvider.closeConnection();
                        // 并重新建立连接和语句对象。
                        establishConnectionAndStatement();
                    }
                } catch (SQLException | ClassNotFoundException exception) {
                    // 如果重新建立连接失败，则抛出运行时异常。
                    LOG.error("JDBC connection is not valid, and reestablish connection failed", exception);
                    throw new RuntimeException("Reestablish JDBC connection failed", exception);
                }
                
                try {
                    // 在每次重试之间，通过 Thread.sleep(1000L * retry)方法暂停一段时间，
                    // 以允许数据库恢复正常状态。
                    Thread.sleep(1000L * retry);
                } catch (InterruptedException e1) {
                    throw new RuntimeException(e1);
                }
            }
        }
        return Collections.emptyList();
    }

    /**
     * 这段代码是在建立数据库连接并准备 SQL语句。
     * 如果在这个过程中出现任何问题（如无法加载数据库驱动，无法建立连接，无法准备语句等），
     * 那么它会抛出 SQLException或 ClassNotFoundException异常。
     * 
     * @throws SQLException
     * @throws ClassNotFoundException
     */
    private void establishConnectionAndStatement() throws SQLException, ClassNotFoundException {
        // 通过调用 connectionProvider的 getOrEstablishConnection方法获取或建立数据库连接，
        // 并将返回的 Connection对象赋值给 dbConn变量。
        Connection dbConn = connectionProvider.getOrEstablishConnection();
        
        //调用 FieldNamedPreparedStatement的 prepareStatement方法，
        // 使用前面获取的数据库连接、查询语句（query）以及键名（keyNames）准备一个预编译的 SQL语句，
        // 并将返回的 PreparedStatement对象赋值给 statement变量。
        statement = FieldNamedPreparedStatement.prepareStatement(dbConn, query, keyNames);
    }

    @Override
    public void close() throws IOException {
        if (statement != null) {
            try {
                statement.close();
            } catch (SQLException e) {
                LOG.info("JDBC statement could not be closed: " + e.getMessage());
            } finally {
                statement = null;
            }
        }
        connectionProvider.closeConnection();
    }

    @VisibleForTesting
    public Connection getDbConnection() {
        return connectionProvider.getConnection();
    }
}
